#PBAR_TestZspecBasicHighZFinder.py
#
#  Testes high-z material finding algorithm. cycles through different windows, 
#
# 10/28/2013, JK, For testing the script, cy

import glob, os
import numpy as np
import PBAR_Zspec, PBAR_Cargo
import matplotlib.pyplot as plt
import copy

# image processing stuff
from skimage.filter import canny
from scipy import ndimage
from skimage import morphology

from skimage.filter import threshold_otsu
from skimage.segmentation import clear_border
from skimage.morphology import label, closing, square
from skimage.measure import regionprops

#from skimage.color import label2rgb
import PIL
from PIL import Image
import pickle
reload(PBAR_Zspec)

figureSize = (16, 10)

# Global parameters

newWidth = 600
newHeight = 136
xCargoInternalBounds = np.array((38, 550))
yCargoInternalBounds = np.array((13, 106))
cargoCountRange = np.array([0, 0.5e8]) # only consider parts of the zspec image that have rad below this value

# make different size windows
windowList = []
windowList.append(np.array((2, 2)))
#windowList.append(np.array((4, 4)))
#windowList.append(np.array((6, 6)))
#windowList.append(np.array((8, 8)))

#discrimThresholdList = np.array((0, 2, 5, 7, 10, 12, 15, 17, 20))
discrimThresholdList = np.array((7, 10, 15))
discrimThresholdList = np.array((7, 10))

#fractionPixelsList = np.array((0.7, 0.8, 0.9, 1.0))
fractionPixelsList = np.array((0.8, 1.0))

(goodZspecMask, badZspecMask, goodZspecNameList, badZspecNameList, \
        goodZspecIndices, badZspecIndices) = PBAR_Zspec.ZspecDetectorLists()

# energy centers
energy = np.array([7.50E+00,1.05E+01,1.35E+01,1.65E+01,1.95E+01,2.25E+01,2.55E+01, \
    2.85E+01,3.15E+01,3.45E+01,3.75E+01,4.05E+01,4.35E+01,4.65E+01,4.95E+01,\
    5.25E+01,5.55E+01,5.85E+01,6.15E+01,6.45E+01,6.75E+01,7.05E+01,7.35E+01, \
    7.65E+01,7.95E+01,8.25E+01,8.55E+01,8.85E+01,9.15E+01,9.45E+01,9.75E+01,1.01E+02])
    
# get linear discriminant coefficients, polynomial fits
(wBest, pfitmean, pfitsigma) = PBAR_Zspec.GetZspecDiscrimValues()

# explore making mask
dataDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScansStandardWidth'
plotDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\HighZFinderPlots'

cargoConfigBaseDirList = glob.glob(os.path.join(dataDir, '*'))

cargoConfigBaseDirList = cargoConfigBaseDirList[35:36]


# round-about way of selecting the desired dataset 
for (dirIndex, cargoConfigBaseDir) in enumerate(cargoConfigBaseDirList):
    subDirList = glob.glob(os.path.join(cargoConfigBaseDir, '*'))
    for (subDirIndex, dataPath) in enumerate(subDirList):        
        print(dataPath)
        a, b = os.path.split(dataPath)
        a, c = os.path.split(a)
        
        plotSaveDir = os.path.join(plotDir, c, b)
        
        # make paths if not present
        if not os.path.exists(os.path.join(plotDir, c)):
            os.mkdir(os.path.join(plotDir, c))

        if not os.path.exists(plotSaveDir):
            os.mkdir(plotSaveDir)        
        
        
        temp = glob.glob(os.path.join(dataPath, '*cargoimage.npy'))
        fullFilenameCargo = temp[0]
        (a, filenameCargo) = os.path.split(fullFilenameCargo)
        # read in the image
        datCargoStandardWidth = np.load(fullFilenameCargo)
        
        # find the zspec scan number
        temp = glob.glob(os.path.join(dataPath, '*.npy'))
        # if file doesn't exist skip this set completely
        a, filename = os.path.split(temp[0])
        zspecScanNumber = filename[0:4]
        filenameZspec = filename
        fullFilenameZspec = os.path.join(dataPath, filename)
        
        # read in the markerfiles
        fullFilenameMarker = os.path.join(dataPath, fullFilenameCargo.replace('cargoimage.npy', 'cargomarker'))
        markerList = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
        
        # calculate discrimination stuff
        datZspecStandardWidth = np.load(fullFilenameZspec)
        discrim = PBAR_Zspec.CalculateFeaturesZspecBasic(datZspecStandardWidth, energy)   
        
#        # Calculate the dist, linear combination for full size image
#        # make features array (num features X time slices X number of detectors)
#        featuresMatrix = np.zeros((3, discrim['binMean'].shape[0], discrim['binMean'].shape[1]))
#        featuresMatrix[0,:,:] = discrim['binMean']
#        featuresMatrix[1,:,:] = discrim['binSTD']
#        featuresMatrix[2,:,:] = discrim['multibin_20_ratio']
#        
#        discrim['dist0']  = featuresMatrix[0,:,:] * wBest[0] + featuresMatrix[1,:,:] * wBest[1] + featuresMatrix[2,:,:] * wBest[2]
#        discrim['discrim'] =  discrim['dist0'] - np.polyval(pfitmean, discrim['count'])
        
        # Reduce the size of the zspec image by summing adjacent bins
        (datZspecSmall, discrimSmall) = PBAR_Zspec.ZspecBasicReduceSize(datZspecStandardWidth, energy, newWidth)
        
        # Modify the marker files
        multiplier = newWidth/float(datZspecStandardWidth.shape[0])
        offset = 0.0
        markerSmallList = PBAR_Zspec.ModifyMarkersXPosition(markerList, multiplier, offset)
        
        ##################
        ##  MAKE MASKs  ##
        
        # Make reduced size cargo image
        datCargoSmall = PBAR_Zspec.CargoReduceSize(datCargoStandardWidth, newHeight, newWidth)
        datCargoSmallMask = PBAR_Zspec.CreateCargoMask(datCargoSmall, cargoCountRange, xCargoInternalBounds, yCargoInternalBounds)
        
        ## Look for high z stuff
        
        # loop over window, discrim threshold and fractionPixels
        for (windowIndex, window) in enumerate(windowList):
            for (discrimThresholdIndex, discrimThreshold) in enumerate(discrimThresholdList):
                for (fractionPixelsIndex, fractionPixels) in enumerate(fractionPixelsList):
                    
                    #
                    print('window: %d, %d; discrim threshold: %d; fractionPixles: %3.2f' %(window[0], window[1], discrimThreshold, fractionPixels))

                    # look for potential stuff
                    potential = PBAR_Zspec.BasicScanZspecHighZFinder(discrimSmall, datZspecSmall, datCargoSmallMask, energy, window, discrimThreshold, fractionPixels)
                        
                    # Save the potential finds to file
                    fullFilenamePotential = \
                        fullFilenameMarker.replace('cargomarker', \
                            '%d_%d_%d_%3.2f.potential' %(window[0], window[1], discrimThreshold, fractionPixels) )
                    with open(fullFilenamePotential, 'wb') as fid:
                        pickle.dump(potential, fid)
                        print('Wrote %s' %fullFilenamePotential)
            
                    ###########
                    ## PLOTS ##
                    ###########
            ;
                    ## Small Cargo image ##
                    
                    plotMarkers = True
                    plotPotential = True
    
                    plt.figure(figsize = figureSize)
                    plt.grid()
                    
                    
                    plt.imshow(datCargoSmall.T, interpolation = 'nearest', aspect='auto', cmap = plt.cm.Greys_r)
                    #plt.imshow(discrimSmall['discrim'].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
                    
                    plt.colorbar()
                    plt.ylabel('Detector Number')
                    plt.xlabel('Time Slice')
                    plt.title('%s, window: %d, %d, disrim threshold: %d, fractionPixels %3.2f' \
                        %(filenameZspec, window[0], window[1], discrimThreshold, fractionPixels))
                    
                    plt.axis((0, 600, 136, 0))
    
                    if plotMarkers:
                        marker = markerSmallList
                        for i, mark in enumerate( marker):
                            x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                            y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                            plt.plot(x, y, 'r')
                            plt.plot(mark['x'], mark['y'], 'xr')
#                            if 'left' in mark:
#                                plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#                            if 'right' in mark:
#                                plt.plot(mark['right']['x'], mark['right']['y'], 'go')
                            plt.text(max(x), min(y), mark['target'], color = 'r', fontsize = 16)
                    
                    if plotPotential:
                        for (potIndex, pot) in enumerate(potential):
                    #        plt.plot(pot['boxCenter'][0], pot['boxCenter'][1], 'xb', markersize = 14)
                            plt.scatter(pot['boxCenter'][0], pot['boxCenter'][1], color = 'g', \
                                s = 100. * pot['window'].mean(), linewidth = pot['window'].mean()*1.0, edgecolor = 'green', facecolor = 'blue', marker = 'x', alpha = 0.5)
                    
    #                plt.show()
                    # save image to file
                    plotFullFilename = os.path.join(plotSaveDir, '%s_%d_%d_%d_%3.2f_Cargo.png' %(zspecScanNumber, window[0], window[1], discrimThreshold, fractionPixels))
                    plt.savefig(plotFullFilename)
                    print('Wrote: %s' %plotFullFilename)
                    plt.close('all')            
            
            
                    # Zspec image, Small with mask
                    
                    climit = (discrimThreshold, 40)
                    
                    plotMarkers = True
                    plotPotential = True

                    plt.figure(figsize = figureSize)
                    plt.grid()
                    
                    mask = ( discrimSmall['count'] > 0 ) & (discrimSmall['count'] < 400) & datCargoSmallMask
                    
                    temp  = copy.copy(discrimSmall['discrim'])
                    cutt = ~np.isnan(temp) & ~np.isinf(temp)
                    temp[~cutt] = min(temp[cutt])
                    
                    temp = discrimSmall['discrim'] * mask
                    minValue = temp[mask].min()
                    temp[~mask] = minValue
                    
                    plt.imshow(temp.T, interpolation = 'nearest', aspect='auto', cmap = plt.cm.Greys_r)
                    #plt.imshow(discrimSmall['discrim'].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
                    
                    plt.colorbar()
                    plt.ylabel('Detector Number')
                    plt.xlabel('Time Slice')
                    plt.title('%s, window: %d, %d, disrim threshold: %d, fractionPixels %3.2f' \
                        %(filenameZspec, window[0], window[1], discrimThreshold, fractionPixels))
                    
                    plt.axis((0, 600, 136, 0))
                    plt.clim(climit)
                    if plotMarkers:
                        marker = markerSmallList
                        for i, mark in enumerate( marker):
                            x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                            y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                            plt.plot(x, y, 'r')
                            plt.plot(mark['x'], mark['y'], 'xr')
#                            if 'left' in mark:
#                                plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
#                            if 'right' in mark:
#                                plt.plot(mark['right']['x'], mark['right']['y'], 'go')
                            plt.text(max(x), min(y), mark['target'], color = 'r', fontsize = 16)
                    
                    if plotPotential:
                        for (potIndex, pot) in enumerate(potential):
                    #        plt.plot(pot['boxCenter'][0], pot['boxCenter'][1], 'xb', markersize = 14)
                            plt.scatter(pot['boxCenter'][0], pot['boxCenter'][1], color = 'g', \
                                s = 100. * pot['window'].mean(), linewidth = pot['window'].mean()*1.0, edgecolor = 'green', facecolor = 'blue', marker = 'x', alpha = 0.5)
                    
                    # save image to file
                    plotFullFilename = os.path.join(plotSaveDir, '%s_%d_%d_%d_%3.2f_ZspecBasic.png' %(zspecScanNumber, window[0], window[1], discrimThreshold, fractionPixels))
                    plt.savefig(plotFullFilename)
                    plt.close('all')
                    
                    
                    # Zspec small, scatter plot
                    removeOffset = True
                    
                    plt.figure(figsize = figureSize)
                    plt.grid()
                    
                    feature1 = 'count'
                    feature2 = 'discrim'
                    
                    mask = (discrimSmall['count'] > 0 ) & (discrimSmall['count'] < 400) & datCargoSmallMask
                    mask = mask[:,goodZspecIndices]
                    
                    x1 = discrimSmall[feature1][:,goodZspecIndices][mask].flatten()
                    y1 = discrimSmall[feature2][:,goodZspecIndices][mask].flatten()

#                    plt.plot(x1, y1, '.x', alpha = 0.10, markersize = 10, label = 'All')
                    
#                    x1 = discrimSmall[feature1][:,goodZspecIndices].flatten()
#                    y1 = discrimSmall[feature2][:,goodZspecIndices].flatten()
                    
                    plt.plot(x1, y1, '.k', alpha = 0.5, markersize = 10, label = 'In Mask')
                    
                    marker = markerSmallList
                    for i, mark in enumerate(marker):
                        x_range = np.array((mark['rec_left'], mark['rec_right'])) # left < right
                        y_range = np.array((mark['rec_top'], mark['rec_bottom'])) # top  < bottom
                        xarray = np.arange(discrimSmall[feature1].shape[0])
                        yarray = np.arange(discrimSmall[feature1].shape[1])
                        xcut = (xarray > x_range[0]) & (xarray < x_range[1])
                        ycut = (yarray > y_range[0]) & (yarray < y_range[1])

                        mask = (discrimSmall['count'] > 0 ) & (discrimSmall['count'] < 400) & datCargoSmallMask
                        mask = mask[:,goodZspecMask & ycut][xcut,:]
                        
                        x = discrimSmall[feature1][:,goodZspecMask & ycut][xcut,:][mask].flatten()
                        y = discrimSmall[feature2][:,goodZspecMask & ycut][xcut,:][mask].flatten()
                    
                        try:
                            # high density stuff
                            if mark['target'][0] == 'S':
                                plt.plot(x, y, 'db', markersize = 8, alpha  = 0.3, label = mark['target'])
                            elif mark['target'][0] == 'W':
                                plt.plot(x, y, 'vb', markersize = 8, alpha  = 0.3, label = mark['target'])
                            elif mark['target'][0] == 'D':
                                plt.plot(x, y, 'sb', markersize = 8, alpha  = 0.3, label = mark['target'])
                            elif mark['target'][0] == 'P':
                                plt.plot(x, y, '*b', markersize = 8, alpha  = 0.3, label = mark['target'])
                            elif mark['target'][0] == 'F':# low density stuff
                                plt.plot(x, y, 'or', markersize = 8, alpha  = 0.3, label = mark['target'])
                            else:
                                plt.plot(x, y, 'om', markersize = 8, alpha  = 0.3, label = mark['target'])
                        except:
                            plt.plot(x, y, 'om', markersize = 8, alpha  = 0.3, label = mark['target'])
                    
                    if plotPotential:
                        for (potIndex, pot) in enumerate(potential):
                            plt.scatter(pot['discrim']['count'], pot['discrim']['discrim'], color = 'g', \
                            s = 100. * pot['window'].mean(), linewidth = pot['window'].mean() * 1.0, edgecolor = 'green', facecolor = 'blue', marker = 'x', alpha = 0.5)
                    
                    # plot threshold line
                    plt.plot([0, 400], [discrimThreshold, discrimThreshold], '-g') 
                    
                    plt.legend()
                    plt.ylabel('Discriminant - Offset')
                    plt.xlabel('Count Per X-ray Pulse')                    
                    
                    plt.title('%s, window: %d, %d, disrim threshold: %d, fractionPixels %3.2f' \
                        %(filenameZspec, window[0], window[1], discrimThreshold, fractionPixels))
                    plt.axis((0, 400, -50, 50))
                
                    # save image to file
                    plotFullFilename = os.path.join(plotSaveDir, '%s_%d_%d_%d_%3.2f_ZspecBasicScatter.png' %(zspecScanNumber, window[0], window[1], discrimThreshold, fractionPixels))
                    plt.savefig(plotFullFilename)
                    plt.close('all')
                    
                    
                        